#pragma once 
#include <crane/aio/buffer.hpp>
#include <crane/core/error.hpp>

#include <boost/asio.hpp>
#include <unordered_map>
#include <stdio.h>
CRA_NS_BEGIN
namespace aio {

class BoostAsio {
public:
    using handle_t = int64_t;
    using AcceptCallback = std::function<void(handle_t ch, std::string ip, int port)>;
    using ConnectCallback = std::function<void(handle_t ch, const std::string & host, const std::string & ip, int port, Error error)>;
    using ReadCallback = std::function<void(size_t bytes, ByteBuffer && buffer, Error error)>;
    using WriteCallback = std::function<void(size_t bytes, ByteBuffer && buffer, Error error)>;

    static const handle_t invalid_handle = -1;

    class SocketManager {
    public:
        using handle_t = int64_t;
        using Result = std::pair<handle_t, boost::asio::ip::tcp::socket *>;
        using SocketMap = std::unordered_map<handle_t, boost::asio::ip::tcp::socket>;
        
        struct Partition {
            SocketMap   sockets;
            size_t      load = 0;
        };

        using PartitionTable = std::vector<Partition>;

        static const handle_t invalid_handle = -1;

    public:
        SocketManager(size_t partitions) : partition_table_(partitions) {}

        Result  add(size_t partition, boost::asio::ip::tcp::socket && socket);
        Result  find(handle_t handle);
        void    erase(handle_t handle);
        size_t  choosePartition() ;                // 获取一个相对空闲的分区，分区负载递增
        void    decreasePartitionLoad(size_t partition);    // 分区负载递减

    private:
        PartitionTable  partition_table_;
    }; // SocketManager

public:
    BoostAsio();

    void stop();
    void run();

    void listen(int port, AcceptCallback && cb);
    void listen(const std::string & ip, int port, AcceptCallback && cb);
    void connect(const std::string & host, int port, ConnectCallback && cb);
    void receiveSome(handle_t ch, ByteBuffer && buffer, ReadCallback && cb);
    void send(handle_t ch, ByteBuffer && buffer, WriteCallback && cb);
    void receive(handle_t ch, ByteBuffer && buffer, ReadCallback && cb);
    void close(handle_t ch);    ///< close channel

private:
    void onAccept(const boost::system::error_code& error, boost::asio::ip::tcp::socket sock);
    void onConnect(
        const boost::system::error_code& error, 
        handle_t ch,
        const std::string& host, const std::string& ip, 
        int port, int partition, ConnectCallback && cb);

private:
    boost::asio::io_context         io_context_;
    std::vector<boost::asio::io_context::strand> strands_;
    boost::asio::ip::tcp::resolver  resolver_;
    boost::asio::ip::tcp::acceptor  acceptor_;
    AcceptCallback                  accept_cb_;
    SocketManager                   sockets_;
}; // BoostAsio

CRA_INLINE
BoostAsio::BoostAsio() 
: resolver_(io_context_)
, acceptor_(io_context_)
, strands_(1, boost::asio::io_context::strand(io_context_))
, sockets_(1) 
{}


CRA_INLINE
void BoostAsio::close(handle_t ch)
{
    int64_t strand_index = (ch >> 32) % strands_.size();
    auto & strand = strands_[strand_index];
    boost::asio::post(strand, [this, ch] () mutable {
        auto result = sockets_.find(ch);
        assert(result.first != invalid_handle);
        
        boost::system::error_code ec;
        result.second->close(ec);

        sockets_.erase(result.first);
    });
}

CRA_INLINE
void  BoostAsio::send(handle_t ch, ByteBuffer && buffer, WriteCallback && cb)
{ 
    int64_t strand_index = (ch >> 32) % strands_.size();
    auto & strand = strands_[strand_index];

    auto fn =  [this, ch, buffer = std::move(buffer), cb = std::move(cb)] () mutable {
        auto result = sockets_.find(ch);
        assert(result.first != invalid_handle);
        boost::asio::const_buffer buf(buffer.data() + buffer.position(), buffer.limit()  - buffer.position());
        boost::asio::async_write(*result.second,
            buf,
            [this, buffer = std::move(buffer), cb = std::move(cb)]( const boost::system::error_code& error, std::size_t bytes_transferred ) mutable
            {
                if ( error ) {
                    Error err = CRA_MAKE_ERROR(error.value(), error.message());
                    cb(bytes_transferred, std::move(buffer), std::move(err));
                }
                else {
                    buffer.setPosition(buffer.position() + bytes_transferred);
                    cb(bytes_transferred, std::move(buffer), Error());
                }
            });
        };

    boost::asio::post(strand, std::move(fn));
}

CRA_INLINE
void  BoostAsio::receiveSome(handle_t ch, ByteBuffer && buffer, ReadCallback && cb)
{ 
    int64_t strand_index = (ch >> 32) % strands_.size();
    auto & strand = strands_[strand_index];
    boost::asio::post(strand, [this, ch, buffer = std::move(buffer), cb = std::move(cb)] () mutable {
        auto result = sockets_.find(ch);
        assert(result.first != invalid_handle);
        boost::asio::mutable_buffer buf(buffer.data() + buffer.position(), buffer.limit() - buffer.position());
        result.second->async_read_some(buf,
            [this, buffer=std::move(buffer), cb = std::move(cb)]( const boost::system::error_code& error, std::size_t bytes_transferred ) mutable
            {
                if ( error ) {
                    Error err = CRA_MAKE_ERROR(error.value(), error.message());
                    cb(bytes_transferred, std::move(buffer), std::move(err));
                }
                else {
                    buffer.setPosition(buffer.position() + bytes_transferred);
                    cb(bytes_transferred, std::move(buffer), Error());
                }
            });
        }
    );
}

CRA_INLINE
void  BoostAsio::receive(handle_t ch, ByteBuffer && buffer, ReadCallback && cb)
{
    int64_t strand_index = (ch >> 32) % strands_.size();
    auto & strand = strands_[strand_index];

    boost::asio::post(
        strand,
        [this, ch, buffer = std::move(buffer), cb = std::move(cb)]() mutable
        {
            auto result = sockets_.find(ch);
            assert(result.first != invalid_handle);
            boost::asio::mutable_buffer buf(buffer.data() + buffer.position(), buffer.limit() - buffer.position());
            boost::asio::async_read(*result.second, buf, 
                [this, buffer = std::move(buffer), cb = std::move(cb)]( const boost::system::error_code& error, std::size_t bytes_transferred ) mutable
                {
                    if ( error ) {
                        Error err = CRA_MAKE_ERROR(error.value(), error.message());
                        cb(bytes_transferred, std::move(buffer), std::move(err));
                    }
                    else {
                        buffer.setPosition(buffer.position() + bytes_transferred);
                        cb(bytes_transferred, std::move(buffer), Error());
                    }
                }
            );
        }
    );
}

CRA_INLINE
void BoostAsio::listen(const std::string & ip, int port, AcceptCallback && cb)
{
    boost::system::error_code ec;
    auto addr = boost::asio::ip::make_address(ip, ec);
    assert(!ec);

    try {
        boost::asio::ip::tcp::acceptor acceptor(io_context_, boost::asio::ip::tcp::endpoint(addr, port));
        this->acceptor_ = std::move(acceptor);
        this->accept_cb_ = std::move(cb);
        this->acceptor_.async_accept(std::bind(&BoostAsio::onAccept, this, std::placeholders::_1, std::placeholders::_2));
    }
    catch (boost::system::system_error & e) {
        abort();
    }
}

CRA_INLINE
void BoostAsio::listen(int port, AcceptCallback && cb)
{
    listen("0.0.0.0", port, std::move(cb));
}

CRA_INLINE
void BoostAsio::connect(const std::string & host, int port, ConnectCallback && cb)
{
    auto fn_resolve = [this, host, port, cb = std::move(cb)] (
            const boost::system::error_code& error, 
            boost::asio::ip::tcp::resolver::results_type results)
    {
        if ( error ) {
            Error err = CRA_MAKE_ERROR(error.value(), error.message());
            cb(invalid_handle, host, "", port, std::move(err));
            return ;   // can not resolve host name
        }

        int32_t partition = this->sockets_.choosePartition();
        boost::asio::post(
            strands_[partition], 
            [this, results=std::move(results), host, port, partition, cb = std::move(cb)] () mutable {
                auto ip = results.begin()->endpoint().address().to_string();
                boost::asio::ip::tcp::socket socket(io_context_);
                socket.open(results.begin()->endpoint().protocol());
                auto result = this->sockets_.add(partition, std::move(socket));
                result.second->async_connect(
                    results.begin()->endpoint(),
                    [this, ch = result.first, host, ip = std::move(ip), port, partition, cb = std::move(cb)] (const boost::system::error_code& error) mutable {
                        this->onConnect(error, ch, host, ip, port, partition, std::move(cb));
                    }
                );
            }
        );
    };

    boost::asio::post(
        strands_[0],
        [this, host, port, fn_resolve = std::move(fn_resolve)] () mutable { 
            resolver_.async_resolve(host, std::to_string(port), std::move(fn_resolve));
        }
    );
}

CRA_INLINE
void BoostAsio::run()
{
    auto work_guard = boost::asio::make_work_guard(io_context_); 
    io_context_.run();
}

CRA_INLINE
void BoostAsio::stop()
{
    io_context_.stop();
}

CRA_INLINE
void BoostAsio::onAccept(const boost::system::error_code& error, boost::asio::ip::tcp::socket sock)
{
    assert(!error);

    std::string ip = sock.remote_endpoint().address().to_string();
    int port = sock.remote_endpoint().port();
    int partition = this->sockets_.choosePartition();   // 获取一个分区

    // 发送到partition对应strand中
    boost::asio::post(strands_[partition], [this, partition, socket = std::move(sock), ip=std::move(ip), port] () mutable {
        auto result = this->sockets_.add(partition, std::move(socket));
        this->accept_cb_(result.first, std::move(ip), port);
    });

    boost::asio::post(io_context_, [this]() {
        this->acceptor_.async_accept(std::bind(&BoostAsio::onAccept, this, std::placeholders::_1, std::placeholders::_2));
    });
}

CRA_INLINE
void BoostAsio::onConnect(
    const boost::system::error_code& error, 
    handle_t ch,
    const std::string& host, 
    const std::string& ip,
    int port, 
    int partition,
    ConnectCallback && cb)
{
    if ( error ) {
        Error err = CRA_MAKE_ERROR(error.value(), error.message());
        this->sockets_.erase(ch);
        cb(invalid_handle, host, ip, port, std::move(err));
        return;
    }

    cb(ch, host, ip, port, CRA_NO_ERROR());
}

CRA_INLINE
void BoostAsio::SocketManager::decreasePartitionLoad(size_t partition)
{
    Partition & part = this->partition_table_[partition];
    part.load--;
}

CRA_INLINE
size_t BoostAsio::SocketManager::choosePartition() 
{
    // TODO 避免连续的choosePartition()返回相同的partition
    size_t partition = 0;
    size_t minimum = INT64_MAX;
    for (size_t i = 0; i < partition_table_.size(); ++i) {
        size_t n = partition_table_[i].load;
        if ( n < minimum ) {
            partition = i;
            minimum = n;
        }
    }
    partition_table_[partition].load++;
    return partition;
}

CRA_INLINE
void BoostAsio::SocketManager::erase(handle_t handle)
{
    int64_t partition_id = handle >> 32;
    fd_t fd = handle & 0xffffffff;

    Partition & partition = partition_table_[partition_id];
    auto it = partition.sockets.find(handle);
    if ( it != partition.sockets.end() ) {
        partition.sockets.erase(it);
        assert(partition.load > 0);
        partition.load -= 1;
    }
}

CRA_INLINE
BoostAsio::SocketManager::Result 
BoostAsio::SocketManager::add(size_t partition, boost::asio::ip::tcp::socket && socket) 
{
    fd_t fd = socket.native_handle();
    handle_t handle = (partition << 32) + fd;
    Partition & part = partition_table_[partition];
    auto r = part.sockets.insert(std::make_pair(handle, std::move(socket)));
    if ( r.second ) {
        return {handle, &r.first->second};
    }
    else {
        return {-1, nullptr};
    }
}

CRA_INLINE
BoostAsio::SocketManager::Result BoostAsio::SocketManager::find(handle_t handle)
{
    int64_t partition = handle >> 32;
    fd_t fd = handle & 0xffffffff;

    Partition & part = partition_table_[partition];
    auto it = part.sockets.find(handle);
    if ( it != part.sockets.end() ) {
        return {handle, &it->second};
    }
    else {
        return {-1, nullptr};
    }
}


} // namespace aio
CRA_NS_END
